import numpy as np 
from matplotlib import pyplot

def euclidean_distance(sample, centers):
    sample = sample.reshape(1, -1)
    centers = centers.reshape(centers.shape[0], -1)
    distances = np.power(np.tile(sample, (centers.shape[0], 1)) - centers, 2).sum(axis=1)
    return distances

class Kmeans(object):
    def __init__(self, K=3, max_iter=500, tolerance=0.0000001):
        self.K_ = K
        self.max_iter_ = max_iter
        self.tolerance_ = tolerance
    def get_closest_centroid(self, sample):
        distances = euclidean_distance(sample, self.centers_)
        closest_centroid = np.argmin(distances)
        return closest_centroid
    def create_clusters(self, X):
        cluster = [[] for _ in range(self.K_)]
        for i, sample in enumerate(X):
            centroid_i = self.get_closest_centroid(sample)
            cluster[centroid_i].append(sample)
        return cluster
    def update_centers(self, clusters, X):
        new_centers = np.zeros((self.K_, np.shape(X)[1]))
        for i, cluster in enumerate(clusters):
            center = np.mean(cluster, axis=0)
            new_centers[i] = center
        return new_centers
    def get_labels(self, clusters, X):
        y_pred = np.zeros(np.shape(X)[0])
        for cluster_i, cluster in enumerate(clusters):
            for sample_i, sample in enumerate(cluster):
                y_pred[sample_i] = cluster_i
        return y_pred
    def fit(self, X, init_centers):
        self.centers_ = init_centers
        # print("centers_: \n", self.centers_)
        for i in range(self.max_iter_):
            
            clusters = self.create_clusters(X)
            # print("clusters: \n", len(clusters))
            new_centers = self.update_centers(clusters, X)
            diff = self.centers_ - new_centers
            if self.tolerance_ == 0:
                self.centers_ = new_centers
            else:
                if abs(diff.any()) < self.tolerance_:
                    break
                self.centers_ = new_centers
            print('epoch ', i + 1)
        return self.get_labels(clusters, X), self.centers_

if __name__ == '__main__':
    centers = np.array([[6.2, 3.2], [6.6, 3.7], [6.5, 3.0]])
    X = np.array([
        [5.9, 3.2],
        [4.6, 2.9],
        [6.2, 2.8],
        [4.7, 3.2],
        [5.5, 4.2],
        [5.0, 3.0],
        [4.9, 3.1],
        [6.7, 3.1],
        [5.1, 3.8],
        [6.0, 3.0]
    ])
    # print(X.shape)
    model = Kmeans(K=3, max_iter=3, tolerance=0)
    result, centers = model.fit(X, centers)
    print(result)
    print('red:')
    print(centers[0])
    print('green:')
    print(centers[1])
    print('blue:')
    print(centers[2])
    

    